﻿#include "stdafx.h"

#include "HookLoadLibrary.h"
#include "HookWinsock.h"

LPCWSTR TARGET_APP  = L"iexplore.exe"; // target application to redirect
LPCSTR  DLL_TOHOOK  = "WS2_32.dll";
LPCWSTR LDLL_TOHOOK = L"WS2_32.dll";

std::wstring GetInstallDir ()
{
    const DWORD len = 256;
    wchar_t     env_var[len];
    
    int ret = ::GetEnvironmentVariableW (L"ProgramFiles", env_var, len);
    if (!ret && ERROR_ENVVAR_NOT_FOUND == GetLastError())
    {
        return std::wstring(L".");
    }
    return std::wstring(env_var);
}

BOOL WINAPI DllMain(HINSTANCE hInst, DWORD Reason, LPVOID reserved)
{
    std::wstring log_file = GetInstallDir();
    log_file += L"\\Socket Redirection for RDS\\DllForInjecting.csv";

    switch(Reason)
    {
    case DLL_PROCESS_ATTACH:
        plog::initialize_log(log_file.c_str(), plog::error);

        // Get name of executable file
        wchar_t NameExe[MAX_PATH+1];
        GetModuleFileName(NULL, NameExe, sizeof(wchar_t)*(MAX_PATH+1));
        wcscpy_s(NameExe, MAX_PATH+1, wcsrchr(NameExe, L'\\') + 1);
        HMODULE hMod;

        if (!lstrcmpi(TARGET_APP, NameExe))
        {
            try
            {
                if(!GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_PIN,LDLL_TOHOOK,&hMod))
                {
                    HookLoadLibrary();
                }
                else
                {
                    HookWinsock2Dll(hMod);
                }
            }
            catch (std::exception& ex)
            {
                LOG(plog::error) << ex.what();
            }
        }
        return TRUE;
    case DLL_PROCESS_DETACH:
        plog::release_log();
        return TRUE;
    case DLL_THREAD_ATTACH:
	case DLL_THREAD_DETACH:
        return TRUE;
    }
    return TRUE;
}
